#include "compi/resource_manager.h"
#include "compi/tag.h"

#include <algorithm>
#include <memory>
#include <mpi.h>
#include <mutex>
#include <set>
#include <stdexcept>
#include <string>
#include <unordered_map>
#include <vector>

namespace compi
{

// 注意：TagAllocator单例现在由ResourceManager管理

/**
 * 命名空间状态管理
 * 每个命名空间维护独立的标签范围和分配状态
 */
struct NamespaceState
{
    std::uint32_t ns_id;
    std::string name;
    std::vector<TagRange> ranges;           // 此命名空间拥有的标签范围
    std::vector<int> free_tags;             // 可重用的标签列表
    std::set<int> allocated_tags;           // 当前已分配的标签集合
    std::unordered_map<int, int> refcounts; // 标签引用计数

    NamespaceState(std::uint32_t id, const std::string & ns_name) : ns_id(id), name(ns_name) {}

    /**
     * 检查标签是否属于此命名空间
     * @param tag_value 标签值
     * @return 如果标签在此命名空间的任一范围内返回true
     */
    bool owns_tag(int tag_value) const
    {
        return std::any_of(ranges.begin(),
                           ranges.end(),
                           [tag_value](const TagRange & range)
                           { return range.contains(tag_value); });
    }

    /**
     * 获取命名空间统计信息
     */
    TagNamespace::Stats get_stats() const
    {
        int total_capacity = 0;
        for (const auto & range : ranges)
        {
            total_capacity += range.size();
        }

        return {static_cast<int>(allocated_tags.size()),
                total_capacity - static_cast<int>(allocated_tags.size()),
                total_capacity,
                static_cast<int>(ranges.size())};
    }
};

struct MessageTagAllocator::Impl
{
    MPI_Comm comm{MPI_COMM_WORLD};
    int tag_ub{32767};
    std::mutex mutex;

    // 命名空间管理
    std::unordered_map<std::uint32_t, std::unique_ptr<NamespaceState>> namespaces;
    std::unordered_map<std::string, std::uint32_t> name_to_id;
    std::uint32_t next_ns_id{1};

    // 全局标签范围管理
    std::set<std::pair<int, int>> allocated_ranges; // (start, end) 已分配的范围
    int next_range_start{1};                        // 下一个可分配范围的起始位置

    // 默认范围大小配置
    static constexpr int DEFAULT_RANGE_SIZE = 256;
    static constexpr int MAX_RANGE_SIZE = 4096;

    /**
     * 为命名空间分配新的标签范围
     * @param ns_state 命名空间状态
     * @param requested_size 请求的范围大小
     * @return 是否成功分配
     */
    bool allocate_range_for_namespace(NamespaceState * ns_state,
                                      int requested_size = DEFAULT_RANGE_SIZE)
    {
        // 确保请求的大小合理
        requested_size = std::min(requested_size, MAX_RANGE_SIZE);
        requested_size = std::max(requested_size, 16);

        // 寻找可用的连续标签范围
        int range_start = find_available_range(requested_size);
        if (range_start < 0 || range_start + requested_size > tag_ub)
        {
            return false; // 无法分配足够的范围
        }

        int range_end = range_start + requested_size;

        // 记录已分配的范围
        allocated_ranges.insert({range_start, range_end});

        // 添加到命名空间
        ns_state->ranges.emplace_back(range_start, range_end);

        // 更新下一个范围起始位置
        next_range_start = range_end;

        return true;
    }

    /**
     * 寻找可用的连续标签范围
     * @param size 需要的范围大小
     * @return 可用范围的起始位置，-1表示无法找到
     */
    int find_available_range(int size)
    {
        int current_start = next_range_start;

        for (const auto & [start, end] : allocated_ranges)
        {
            if (current_start + size <= start)
            {
                // 在当前位置和下一个已分配范围之间有足够空间
                return current_start;
            }
            current_start = std::max(current_start, end);
        }

        // 检查最后一个范围之后是否有足够空间
        if (current_start + size <= tag_ub)
        {
            return current_start;
        }

        return -1; // 无法找到足够的空间
    }

    /**
     * 从命名空间分配标签
     * @param ns_state 命名空间状态
     * @return 分配的标签值，-1表示失败
     */
    int allocate_tag_from_namespace(NamespaceState * ns_state)
    {
        // 1. 优先使用空闲标签列表
        if (!ns_state->free_tags.empty())
        {
            int tag = ns_state->free_tags.back();
            ns_state->free_tags.pop_back();
            ns_state->allocated_tags.insert(tag);
            ns_state->refcounts[tag] = 1;
            return tag;
        }

        // 2. 从现有范围中分配新标签
        for (const auto & range : ns_state->ranges)
        {
            for (int tag = range.start; tag < range.end; ++tag)
            {
                if (ns_state->allocated_tags.find(tag) == ns_state->allocated_tags.end())
                {
                    ns_state->allocated_tags.insert(tag);
                    ns_state->refcounts[tag] = 1;
                    return tag;
                }
            }
        }

        // 3. 尝试为命名空间分配新范围
        if (allocate_range_for_namespace(ns_state))
        {
            // 从新分配的范围中取第一个标签
            const auto & new_range = ns_state->ranges.back();
            int tag = new_range.start;
            ns_state->allocated_tags.insert(tag);
            ns_state->refcounts[tag] = 1;
            return tag;
        }

        return -1; // 分配失败
    }

    /**
     * 释放标签到命名空间
     * @param ns_state 命名空间状态
     * @param tag_value 要释放的标签值
     */
    void release_tag_to_namespace(NamespaceState * ns_state, int tag_value)
    {
        auto ref_it = ns_state->refcounts.find(tag_value);
        if (ref_it == ns_state->refcounts.end())
        {
            return; // 标签不存在
        }

        ref_it->second--;
        if (ref_it->second <= 0)
        {
            // 引用计数归零，释放标签
            ns_state->refcounts.erase(ref_it);
            ns_state->allocated_tags.erase(tag_value);
            ns_state->free_tags.push_back(tag_value);
        }
    }
};

static int
query_tag_ub(MPI_Comm comm)
{
    int tag_ub = 32767; // MPI标准最小值
    int flag = 0;
    void * attr_val = nullptr;

    if (MPI_Comm_get_attr(comm, MPI_TAG_UB, &attr_val, &flag) == MPI_SUCCESS && flag)
    {
        tag_ub = *static_cast<int *>(attr_val);
    }

    return tag_ub;
}

MessageTagAllocator::MessageTagAllocator(MPI_Comm comm) : _impl(std::make_unique<Impl>())
{
    _impl->comm = comm;
    _impl->tag_ub = query_tag_ub(comm);
}

MessageTagAllocator::~MessageTagAllocator() = default;

TagNamespace
MessageTagAllocator::create_namespace(const std::string & name)
{
    std::lock_guard<std::mutex> lock(_impl->mutex);

    // 检查名称是否已存在
    if (!name.empty() && _impl->name_to_id.find(name) != _impl->name_to_id.end())
    {
        throw std::runtime_error("Namespace name already exists: " + name);
    }

    std::uint32_t ns_id = _impl->next_ns_id++;
    auto ns_state = std::make_unique<NamespaceState>(ns_id, name);

    // 为新命名空间分配初始范围
    if (!_impl->allocate_range_for_namespace(ns_state.get()))
    {
        throw std::runtime_error("Failed to allocate initial range for namespace");
    }

    // 注册命名空间
    _impl->namespaces[ns_id] = std::move(ns_state);
    if (!name.empty())
    {
        _impl->name_to_id[name] = ns_id;
    }

    return TagNamespace(ns_id, this);
}

TagNamespace
MessageTagAllocator::find_namespace(const std::string & name)
{
    std::lock_guard<std::mutex> lock(_impl->mutex);

    auto it = _impl->name_to_id.find(name);
    if (it == _impl->name_to_id.end())
    {
        throw std::runtime_error("Namespace not found: " + name);
    }

    return TagNamespace(it->second, this);
}

void
MessageTagAllocator::destroy_namespace(std::uint32_t ns_id)
{
    std::lock_guard<std::mutex> lock(_impl->mutex);

    auto it = _impl->namespaces.find(ns_id);
    if (it == _impl->namespaces.end())
    {
        return; // 命名空间不存在
    }

    // 释放命名空间的所有范围
    for (const auto & range : it->second->ranges)
    {
        _impl->allocated_ranges.erase({range.start, range.end});
    }

    // 从名称映射中移除
    if (!it->second->name.empty())
    {
        _impl->name_to_id.erase(it->second->name);
    }

    // 移除命名空间
    _impl->namespaces.erase(it);
}

int
MessageTagAllocator::tag_ub() const noexcept
{
    return _impl->tag_ub;
}

MessageTagAllocator::GlobalStats
MessageTagAllocator::get_global_stats() const
{
    std::lock_guard<std::mutex> lock(_impl->mutex);

    int total_allocated = 0;
    int total_available = 0;

    for (const auto & [ns_id, ns_state] : _impl->namespaces)
    {
        auto stats = ns_state->get_stats();
        total_allocated += stats.allocated_count;
        total_available += stats.available_count;
    }

    return {static_cast<int>(_impl->namespaces.size()),
            total_allocated,
            total_available,
            _impl->tag_ub};
}

MessageTag
MessageTagAllocator::request(std::uint32_t ns_id)
{
    std::lock_guard<std::mutex> lock(_impl->mutex);

    auto it = _impl->namespaces.find(ns_id);
    if (it == _impl->namespaces.end())
    {
        throw std::runtime_error("Namespace not found: " + std::to_string(ns_id));
    }

    int tag_value = _impl->allocate_tag_from_namespace(it->second.get());
    if (tag_value < 0)
    {
        throw std::runtime_error("Failed to allocate tag for namespace: " + std::to_string(ns_id));
    }

    return MessageTag(tag_value, ns_id);
}

void
MessageTagAllocator::retain(const MessageTag & tag)
{
    std::lock_guard<std::mutex> lock(_impl->mutex);

    auto it = _impl->namespaces.find(tag.ns_id);
    if (it != _impl->namespaces.end())
    {
        auto & refcounts = it->second->refcounts;
        auto ref_it = refcounts.find(tag.value);
        if (ref_it != refcounts.end())
        {
            ref_it->second++;
        }
    }
}

void
MessageTagAllocator::release(const MessageTag & tag)
{
    std::lock_guard<std::mutex> lock(_impl->mutex);

    auto it = _impl->namespaces.find(tag.ns_id);
    if (it != _impl->namespaces.end())
    {
        _impl->release_tag_to_namespace(it->second.get(), tag.value);
    }
}

bool
MessageTagAllocator::belongs(std::uint32_t ns_id, int tag_value) const
{
    std::lock_guard<std::mutex> lock(_impl->mutex);

    auto it = _impl->namespaces.find(ns_id);
    if (it == _impl->namespaces.end())
    {
        return false;
    }

    return it->second->owns_tag(tag_value);
}

// TagNamespace 实现

TagNamespace::TagNamespace(std::uint32_t ns_id, MessageTagAllocator * allocator) : _ns_id(ns_id)
{
    // 将裸指针包装为 shared_ptr，但不管理生命周期（空删除器）
    _owner = std::shared_ptr<MessageTagAllocator>(allocator, [](MessageTagAllocator *) {});
}

std::uint32_t
TagNamespace::id() const noexcept
{
    return _ns_id;
}

const std::string &
TagNamespace::name() const noexcept
{
    std::lock_guard<std::mutex> lock(_owner->_impl->mutex);

    auto it = _owner->_impl->namespaces.find(_ns_id);
    if (it != _owner->_impl->namespaces.end())
    {
        return it->second->name;
    }

    static const std::string empty_name;
    return empty_name;
}

MessageTag
TagNamespace::request() const
{
    return _owner->request(_ns_id);
}

void
TagNamespace::retain(const MessageTag & tag) const
{
    _owner->retain(tag);
}

void
TagNamespace::release(const MessageTag & tag) const
{
    _owner->release(tag);
}

bool
TagNamespace::owns(int tag_value) const
{
    return _owner->belongs(_ns_id, tag_value);
}

std::vector<TagRange>
TagNamespace::get_ranges() const
{
    std::lock_guard<std::mutex> lock(_owner->_impl->mutex);

    auto it = _owner->_impl->namespaces.find(_ns_id);
    if (it != _owner->_impl->namespaces.end())
    {
        return it->second->ranges;
    }

    return {};
}

TagNamespace::Stats
TagNamespace::get_stats() const
{
    std::lock_guard<std::mutex> lock(_owner->_impl->mutex);

    auto it = _owner->_impl->namespaces.find(_ns_id);
    if (it != _owner->_impl->namespaces.end())
    {
        return it->second->get_stats();
    }

    return {0, 0, 0, 0};
}

bool
TagNamespace::iprobe(MPI_Comm comm, int src_rank, MessageTag & out_tag) const
{
    // 获取命名空间的所有范围
    auto ranges = get_ranges();

    for (const auto & range : ranges)
    {
        for (int tag = range.start; tag < range.end; ++tag)
        {
            int flag = 0;
            MPI_Status status;

            int result = MPI_Iprobe(src_rank, tag, comm, &flag, &status);
            if (result == MPI_SUCCESS && flag)
            {
                out_tag = MessageTag(tag, _ns_id);
                return true;
            }
        }
    }

    return false;
}

TagLease
TagNamespace::lease() const
{
    auto tag = _owner->request(_ns_id);
    return TagLease(_owner, _ns_id, tag, true);
}

// TagLease 实现

TagLease::~TagLease() { retire(); }

TagLease::TagLease(TagLease && other) noexcept
    : _owner(std::move(other._owner)),
      _ns_id(other._ns_id),
      _tag(other._tag),
      _active(other._active)
{
    other._active = false;
    other._ns_id = 0;
    other._tag = MessageTag(0, 0);
}

TagLease &
TagLease::operator=(TagLease && other) noexcept
{
    if (this == &other)
        return *this;
    // 先释放当前持有
    retire();
    _owner = std::move(other._owner);
    _ns_id = other._ns_id;
    _tag = other._tag;
    _active = other._active;
    other._active = false;
    other._ns_id = 0;
    other._tag = MessageTag(0, 0);
    return *this;
}

void
TagLease::retire() noexcept
{
    if (_active && _owner)
    {
        try
        {
            _owner->release(_tag);
        }
        catch (...)
        {
            // 析构路径下吞掉异常
        }
        _active = false;
    }
}

bool
TagLease::renew()
{
    if (!_owner)
        return false;
    // 若当前活跃，先归还旧标签
    if (_active)
    {
        try
        {
            _owner->release(_tag);
        }
        catch (...)
        {
        }
    }
    try
    {
        _tag = _owner->request(_ns_id);
        _active = true;
        return true;
    }
    catch (...)
    {
        _active = false;
        return false;
    }
}

} // namespace compi